#!/bin/bash
# Copyright (C) 2024-2025 KoudaiPC Computer Technologies Co., Ltd.
# script name: yandi-host-setup-nocloud-amd64-image.sh
# author: david921518@qq.com

################
# global vars
GLB_WORK_DIR=
GLB_BOOTSTRAP_FILE=
GLB_BOOTSTRAP_ROOT=
GLB_PREINSTALL_ROOT=
GLB_BOOTSTRAP_DEBS=
GLB_PREINSTALL_LIST=
################

user=$(whoami)

# 定义日志格式-前缀
log_prefix="[yandi-setup]"
# 定义日志目录和文件名
log_file_dir="/var/log/kpc/yandi"
log_file_path="$log_file_dir/host-setup.log"

log_kmsg_path="/dev/kmsg"
log_con_path="/dev/tty1"

function log_init()
{
  # 文件不存在则创建
  if [ ! -d "$log_file_dir" ]; then
    mkdir -p $log_file_dir
    chmod -R 755 $log_file_dir
  fi
  if [ ! -f "$log_file_path" ]; then
    touch $log_file_path
  fi
}

# 定义日志打印格式
function log()
{
  # 定义日志时间格式
  local curtime=$(date "+%Y-%m-%d %H:%M:%S")
  echo "$curtime $log_prefix $@"
  if [ "$user" == "root" ]; then
    echo "$curtime $log_prefix $@" >>$log_file_path
    #sleep 0.1
    #echo "$log_prefix [kmsg] $@" >$log_kmsg_path
    #sleep 0.1
    #echo "$curtime $log_prefix [con] $@" >$log_con_path
    #sleep 0.1
  fi
}

function log_from_file()
{
  tmpfile=$@
  if [ ! -f "$tmpfile" ]; then
    log "no file [$tmpfile] to log"
    return 1
  fi

  n=0
  while read line;
  do
    lists[$n]="$line"
    ((n++))
  done <$tmpfile

  # 定义日志时间格式
  local curtime=$(date "+%Y-%m-%d %H:%M:%S")

  #echo "$curtime $log_prefix $@"
  for i in $( seq 0 $((${#lists[*]} - 1)) )
  do
    echo "$curtime $log_prefix ${lists[i]}"
  done

  if [ "$user" == "root" ]; then
    #echo "$curtime $log_prefix $@" >>$log_file_path
    for i in $( seq 0 $((${#lists[*]} - 1)) )
    do
      echo "$curtime $log_prefix ${lists[i]}" >>$log_file_path
    done

    #sleep 0.1

    #echo "$log_prefix [kmsg] $@" >$log_kmsg_path
    #for i in $( seq 0 $((${#lists[*]} - 1)) )
    #do
    #  echo "$curtime $log_prefix ${lists[i]}" >$log_kmsg_path
    #  [ $((i%10)) -ne 9 ] || sleep 0.1
    #done

    #sleep 0.1

    #echo "$curtime $log_prefix [con] $@" >$log_con_path
    #for i in $( seq 0 $((${#lists[*]} - 1)) )
    #do
    #  echo "$curtime $log_prefix ${lists[i]}" >$log_con_path
    #done

    #sleep 0.1
  fi
}

function get_work_dir()
{
  ################
  # /kpcos/distro/yandi/releases/1.0/ports/amd64/host/yandi-host-setup-nocloud-amd64-image.sh
  # /kpcos/distro/yandi/bootstrap.sh
  ################

  local real_path=$(realpath $0)
  GLB_WORK_DIR=$(dirname "$real_path")
  log "WORK DIR is [$GLB_WORK_DIR]"

  local bootstrap_file="$GLB_WORK_DIR/../../../../../bootstrap.sh"
  log "bootstrap_file=[$bootstrap_file]"
  real_bootstrap_file=$(realpath "$bootstrap_file")
  log "real_bootstrap_file=[$real_bootstrap_file]"

  if [ ! -f $real_bootstrap_file ]; then
    log "No bootstrap [$real_bootstrap_file]."
    exit 1
  fi

  if [ ! -x $real_bootstrap_file ]; then
    log "bootstrap [$real_bootstrap_file] is not executive."
    exit 1
  fi

  GLB_BOOTSTRAP_FILE="$real_bootstrap_file"
  log "GLB_BOOTSTRAP_FILE=[$GLB_BOOTSTRAP_FILE]"

  local bootstrap_root="$GLB_WORK_DIR/../bootstrap"
  log "bootstrap_root=[$bootstrap_root]"
  real_bootstrap_root=$(realpath "$bootstrap_root")
  log "real_bootstrap_root=[$real_bootstrap_root]"

  if [ ! -d $real_bootstrap_root ]; then
    log "No bootstrap root dir [$real_bootstrap_root]."
    exit 1
  fi

  GLB_BOOTSTRAP_ROOT="$real_bootstrap_root"
  log "GLB_BOOTSTRAP_ROOT=[$GLB_BOOTSTRAP_ROOT]"

  GLB_PREINSTALL_ROOT="$GLB_BOOTSTRAP_ROOT/preinstall"
  log "GLB_PREINSTALL_ROOT=[$GLB_PREINSTALL_ROOT]"

  GLB_BOOTSTRAP_DEBS="$GLB_PREINSTALL_ROOT/bootstrap_debs.txt"
  log "GLB_BOOTSTRAP_DEBS=[$GLB_BOOTSTRAP_DEBS]"

  GLB_PREINSTALL_LIST="$GLB_PREINSTALL_ROOT/preinstall_debs.txt"
  log "GLB_PREINSTALL_LIST=[$GLB_PREINSTALL_LIST]"
}

# check WAN interface
function check_wan_iface()
{
  let COUNTER=0
  iface_dev=$(ip route | awk '/default/ {print $5}')
  while [ "${iface_dev}"x == x ]; do
    sleep 1
    let COUNTER=$COUNTER+1
    if [ $COUNTER -gt 29 ]; then
      log "Can't get default route device interface, exit now"
      exit 1
    fi
    log "get default route device interface again"
    iface_dev=$(ip route | awk '/default/ {print $5}')
  done
  let COUNTER=0
  WAN_IP=$(ip addr show $iface_dev | grep "inet" | head -n 1 | awk '/inet/ {print $2}' | cut -d'/' -f1)
  while [ "${WAN_IP}"x == x ]; do
    sleep 1
    let COUNTER=$COUNTER+1
    if [ $COUNTER -gt 29 ]; then
      log "Can't get WAN IP address, exit now"
      exit 1
    fi
    log "get WAN IP address again"
    WAN_IP=$(ip addr show $iface_dev | grep "inet" | head -n 1 | awk '/inet/ {print $2}' | cut -d'/' -f1)
  done
  log "WAN IP=[${WAN_IP}]"
}

# get bootstrap deb packages
function get_bootstrap_debs()
{
  if [ ! -f $GLB_BOOTSTRAP_DEBS ]; then
    log "no file $GLB_BOOTSTRAP_DEBS , return now"
    return 1
  fi

  n=0
  while read line;
  do
    lists[$n]="$line"
    ((n++))
  done <$GLB_BOOTSTRAP_DEBS

  for i in $( seq 0 $((${#lists[*]} - 1)) )
  do
    #log "lists[$i]=[${lists[i]}]"
    name=$(echo ${lists[i]} | awk '{print $1}')
    url=$(echo ${lists[i]} | awk '{print $2}')
    log "name=[$name]"
    log "url=[$url]"
    case "$name" in
      "#"* )
        #log "it's comment"
        ;;
      "" )
        #log "it's empty line"
        ;;
      * )
        #log "name=[$name]"
	filename="${url##*/}"
        if [ "${filename}" != "" ]; then
	  md5_file="$GLB_PREINSTALL_ROOT"/${name}.md5

          need_dl="yes"
          md5=""

	  if [ -f $md5_file ]; then
            md5=$(cat $md5_file)
	  fi
          log "stored md5=[$md5]"

          # check if $filename exist
          if [ -f $filename ]; then
            MD5SUM=$(md5sum $filename | awk '{print $1}')
            if [ "$MD5SUM" != "$md5" ]; then
              log "File $filename is incorrect, download again!"
              rm -f $filename
              need_dl="yes"
            else
              log "File $filename is correct"
              need_dl="no"
            fi
          fi

	  if [ "$need_dl" != "no" ]; then
            log "curl -o $filename -L $url"
            curl -o $filename -L $url
          fi

	  if [ ! -f $filename ]; then
            echo "No file $filename"
            exit 1
          fi

	  if [ ! -f $md5_file ]; then
            MD5SUM=$(md5sum $filename | awk '{print $1}')
	    echo "$MD5SUM" >${name}.md5

            log "cp -f ${name}.md5 $md5_file"
            cp -f ${name}.md5 $md5_file
	  fi
          md5=$(cat $md5_file)

          MD5SUM=$(md5sum $filename | awk '{print $1}')
          if [ "$MD5SUM" != "$md5" ]; then
            echo "File $filename is incorrect, md5_calc=[$MD5SUM] md5=[$md5]"
            exit 1
          fi

	  log "cp -f $filename $GLB_PREINSTALL_ROOT/"
	  cp -f $filename $GLB_PREINSTALL_ROOT/
        fi
        ;;
    esac
  done
}

# get preinstall packages
function get_preinstall_pkgs()
{
  if [ ! -f $GLB_PREINSTALL_LIST ]; then
    log "no file $GLB_PREINSTALL_LIST , return now"
    return 1
  fi

  n=0
  while read line;
  do
    lists[$n]="$line"
    ((n++))
  done <$GLB_PREINSTALL_LIST

  for i in $( seq 0 $((${#lists[*]} - 1)) )
  do
    #log "lists[$i]=[${lists[i]}]"
    name=$(echo ${lists[i]} | awk '{print $1}')
    url=$(echo ${lists[i]} | awk '{print $2}')
    log "name=[$name]"
    log "url=[$url]"
    case "$name" in
      "#"* )
        #log "it's comment"
        ;;
      "" )
        #log "it's empty line"
        ;;
      * )
        #log "name=[$name]"
	filename="${url##*/}"
        if [ "${filename}" != "" ]; then
	  md5_file="$GLB_PREINSTALL_ROOT"/${name}.md5

          need_dl="yes"
          md5=""

	  if [ -f $md5_file ]; then
            md5=$(cat $md5_file)
	  fi
          log "stored md5=[$md5]"

          # check if $filename exist
          if [ -f $filename ]; then
            MD5SUM=$(md5sum $filename | awk '{print $1}')
            if [ "$MD5SUM" != "$md5" ]; then
              log "File $filename is incorrect, download again!"
              rm -f $filename
              need_dl="yes"
            else
              log "File $filename is correct"
              need_dl="no"
            fi
          fi

	  if [ "$need_dl" != "no" ]; then
            log "curl -o $filename -L $url"
            curl -o $filename -L $url
          fi

	  if [ ! -f $filename ]; then
            echo "No file $filename"
            exit 1
          fi

	  if [ ! -f $md5_file ]; then
            MD5SUM=$(md5sum $filename | awk '{print $1}')
	    echo "$MD5SUM" >${name}.md5

            log "cp -f ${name}.md5 $md5_file"
            cp -f ${name}.md5 $md5_file
	  fi
          md5=$(cat $md5_file)

          MD5SUM=$(md5sum $filename | awk '{print $1}')
          if [ "$MD5SUM" != "$md5" ]; then
            echo "File $filename is incorrect, md5_calc=[$MD5SUM] md5=[$md5]"
            exit 1
          fi

	  log "cp -f $filename $GLB_PREINSTALL_ROOT/"
	  cp -f $filename $GLB_PREINSTALL_ROOT/
        fi
        ;;
    esac
  done
}

function setup_image()
{
  local img_file=$1
  local type=$2
  local type_bootstrap_conf="bootstrap.conf.qemu"

  if [ ! -f "$img_file" ]; then
    log "no image file [$img_file] to setup"
    return 1
  fi

  if [ "$type" == "udisk" ]; then
    type_bootstrap_conf="bootstrap.conf.udisk"
  fi

  #nbd_dev=/dev/nbd0
  nbd_dev=
  for x in /sys/class/block/nbd* ; do
    S=$(cat $x/size)
    if [ "$S" == "0" ] ; then
      nbd_dev="/dev/$(basename $x)"
      break
    fi
  done
  if [ x"$nbd_dev" == x ]; then
    log "can't find free NBD device"
    return 1
  fi
  log "nbd_dev=[$nbd_dev]"

  # 2. Connect the QCOW2 as a network block device
  #qemu-nbd --connect=/dev/nbd0 ./kpcos-yandi-1.0-nocloud-amd64-40g.qcow2
  qemu-nbd --connect=$nbd_dev $img_file

  let counter=0
  while [ ! -e "${nbd_dev}p1" ]
  do
    if [ $counter -gt 9 ]; then
      log "wait too long for ${nbd_dev}p1 being ready, exit now."
      qemu-nbd --disconnect $nbd_dev
      exit 1
    fi
    let counter=$counter+1
    log "${nbd_dev}p1 is not ready, wait 1 second."
    sleep 1
  done

  # 3. Fix qcow2 image gpt
  #rm -f /tmp/parted_info
  #/sbin/parted --script /dev/nbd0 p 2>&1 | tee -a /tmp/parted_info
  # gpt warning
  #if grep -Eq 'fix the GPT' /tmp/parted_info; then
  #  echo -e "Fix\n" | /sbin/parted ---pretend-input-tty /dev/nbd0 print 1>/dev/null
  #  rm -f /tmp/parted_info
  #fi
  info_file=$(mktemp)
  /sbin/parted --script $nbd_dev print 2>&1 | tee $info_file
  log_from_file $info_file
  while grep -Eq 'Error:' $info_file
  do
    echo -e "Ignore\nOK\nFix\n" | /sbin/parted ---pretend-input-tty $nbd_dev rescue 1>/dev/null
    /sbin/parted --script $nbd_dev print 2>&1 | tee $info_file
    log_from_file $info_file
    sleep 9
  done

  if grep -Eq 'fix the GPT' $info_file ; then
    echo -e "Fix\n" | /sbin/parted ---pretend-input-tty $nbd_dev print 1>/dev/null
  fi
  rm -f $info_file

  # 4. enlarge partition
  log "parted $nbd_dev resizepart 1 100%"
  tmpfile=$(mktemp)
  /sbin/parted $nbd_dev resizepart 1 100% 2>&1 | tee $tmpfile
  log_from_file $tmpfile
  rm -f $tmpfile

  log "parted $nbd_dev print"
  tmpfile=$(mktemp)
  /sbin/parted $nbd_dev print 2>&1 | tee $tmpfile
  log_from_file $tmpfile
  rm -f $tmpfile

  # 5. enlarge e2fs partition
  loop_device=$(/sbin/losetup -f)
  nbd_dev_p1="${nbd_dev}p1"
  mount_point=/mnt/yandi
  if [ ! -d $mount_point ]; then
    mkdir -p $mount_point
  fi
  /sbin/losetup $loop_device $nbd_dev_p1
  sleep 0.1

  log "e2fsck -fy $loop_device"
  tmpfile=$(mktemp)
  /sbin/e2fsck -fy $loop_device 2>&1 | tee $tmpfile
  log_from_file $tmpfile
  rm -f $tmpfile
  sleep 0.1

  log "resize2fs -f $loop_device"
  tmpfile=$(mktemp)
  /sbin/resize2fs -f $loop_device 2>&1 | tee $tmpfile
  log_from_file $tmpfile
  rm -f $tmpfile
  sleep 0.1

  /sbin/losetup -d $loop_device
  sleep 0.1

  # 6. Mount the partition from the VM
  log "mount $nbd_dev_p1 $mount_point"
  mount $nbd_dev_p1 $mount_point
  sleep 0.1
  let counter=0
  while [ ! -e "$mount_point/etc" ]
  do
    if [ $counter -gt 9 ]; then
      log "wait too long for $mount_point/etc being ready, exit now."
      qemu-nbd --disconnect $nbd_dev
      exit 1
    fi
    let counter=$counter+1
    log "$mount_point/etc is not ready, wait 1 second."
    sleep 1
  done

  # 7. Copy target scripts to mounted point
  #####
  # only copy bootstrap needed files
  log "cp -af $GLB_BOOTSTRAP_ROOT/etc/* $mount_point/etc/"
  cp -af "$GLB_BOOTSTRAP_ROOT"/etc/* $mount_point/etc/

  log "cp -af $GLB_BOOTSTRAP_ROOT/lib/systemd/* $mount_point/lib/systemd/"
  cp -af "$GLB_BOOTSTRAP_ROOT"/lib/systemd/* $mount_point/lib/systemd/

  log "cp -af $GLB_BOOTSTRAP_ROOT/usr/sbin/* $mount_point/usr/sbin/"
  cp -af "$GLB_BOOTSTRAP_ROOT"/usr/sbin/* $mount_point/usr/sbin/

  #mkdir -p /mnt/yandi/root
  log "mkdir -p $mount_point/etc/kpc"
  mkdir -p $mount_point/etc/kpc
  log "cp -af $GLB_BOOTSTRAP_FILE $mount_point/etc/kpc/"
  cp -af $GLB_BOOTSTRAP_FILE $mount_point/etc/kpc/

  log "mkdir -p $mount_point/etc/kpc/bootstrap"
  mkdir -p $mount_point/etc/kpc/bootstrap
  log "cp -af $GLB_BOOTSTRAP_ROOT/* $mount_point/etc/kpc/bootstrap/"
  cp -af "$GLB_BOOTSTRAP_ROOT"/* $mount_point/etc/kpc/bootstrap/

  # install type bootstrap conf
  log "cp -af $GLB_BOOTSTRAP_ROOT/etc/kpc/$type_bootstrap_conf $mount_point/etc/kpc/bootstrap.conf"
  cp -af "$GLB_BOOTSTRAP_ROOT"/etc/kpc/$type_bootstrap_conf $mount_point/etc/kpc/bootstrap.conf

  log "chown root:root $mount_point/root -R"
  chown root:root $mount_point/root -R
  log "chown root:root $mount_point/etc -R"
  chown root:root $mount_point/etc -R
  log "chown root:root $mount_point/lib/systemd -R"
  chown root:root $mount_point/lib/systemd -R
  log "chown root:root $mount_point/usr/sbin -R"
  chown root:root $mount_point/usr/sbin -R
  #####
  #mkdir -p /mnt/yandi/etc/kpc/firstboot
  #cp -af ../guest /mnt/yandi/etc/kpc/firstboot/
  #cp -af ../root /mnt/yandi/etc/kpc/firstboot/
  #cp -af ../root/etc/* /mnt/yandi/etc/
  #cp -af ../root/lib/systemd/* /mnt/yandi/lib/systemd/
  #cp -af ../root/usr/sbin/* /mnt/yandi/usr/sbin/
  #chown root:root /mnt/yandi/etc -R
  #chown root:root /mnt/yandi/lib/systemd -R
  #chown root:root /mnt/yandi/usr/sbin -R
  #chown root:root /mnt/yandi/etc/kpc/firstboot -R
  #####

  # 8. Unmount
  log "umount $mount_point"
  umount $mount_point
  sleep 0.1

  # 9. Disconnect
  log "qemu-nbd --disconnect $nbd_dev"
  qemu-nbd --disconnect $nbd_dev
}

# 0. init scripts
log_init
get_work_dir
cd "$GLB_WORK_DIR"
check_wan_iface
get_bootstrap_debs
get_preinstall_pkgs

# 1. Enable NBD on the host
/sbin/modprobe nbd max_part=8

size_list=(
  "3g" "40g"
)

for size in "${size_list[@]}"
do
  if [ "$size" == "3g" ]; then
    install_type="udisk"
  else
    install_type="qemu"
  fi
  log "install_type=[$install_type]"
  img_file="kpcos-yandi-1.02-nocloud-amd64-${size}.qcow2"
  log "img_file=[$img_file]"
  if [ -f $img_file ]; then
    log "setup image file [$img_file]"
    setup_image $img_file $install_type
  fi
done

exit 0
